from sklearn.datasets import load_iris
import matplotlib.pyplot as plt
import numpy as np
from sklearn import metrics

# 数据获取
def get_IRIS_data():
    iris = load_iris()
    data = iris.data
    label = iris.target
    return data, label

# k-means初始化类中心
def select_centre(data, k):
    center_list = []
    for _ in range(k):
        ran = np.random.randint(len(data))
        center_list.append(data[ran])
    return center_list

# k-means聚类
def k_means(data, label, center, iter):
    for _ in range(iter):
        label_result = []
        data_result = []
        clusterAssment = np.zeros([len(data), 2])
        label_pre = []
        
        for img, i in zip(data, range(len(data))):
            minDis = 1e9
            minIndex = 0
            for cen, j in zip(center, range(len(center))):
                # 欧氏距离进行距离度量
                dis = np.linalg.norm(img - cen)
                if dis < minDis:
                    minDis = dis
                    minIndex = j
            clusterAssment[i, :] = minIndex, minDis
            label_pre.append(minIndex)

        for j in range(len(center)):
            pointsInCluster = data[clusterAssment[:, 0] == j]
            labelInCluster = label[clusterAssment[:, 0] == j]
            center[j] = np.mean(pointsInCluster, axis=0)
            label_result.append(labelInCluster)
            data_result.append(pointsInCluster)

    return label_pre, label_result, data_result, center

# 引入兰德指数评价聚类结果
def metric_ARI(result, pred_result):
    return metrics.adjusted_rand_score(result, pred_result)

if __name__ == '__main__':
    # 加载数据集
    data, label = get_IRIS_data()
    # 初始化类中心
    centers = select_centre(data, 3)
    # k-means
    label_pre, label_result, data_result, new_centers = k_means(data, label, centers, iter = 10)
    # 评价指标
    print('兰德指数: ' + str(metric_ARI(label, label_pre)))



